import sys

import numpy as np
from PyQt5 import QtGui, QtCore
from PyQt5.QtCore import pyqtSignal, Qt, QSize
from PyQt5.QtGui import QPalette
from PyQt5.QtWidgets import QApplication, QWidget
from matplotlib.backends.backend_qt5agg import FigureCanvasQTAgg
from matplotlib.figure import Figure
from scipy import interpolate
from ui.CurveWidget import Ui_Curve


def normalize(pos_x, pos_y):
    x = (pos_x - 50) * (255 / 306)
    y = 255 - (pos_y - 50) * (255 / 306)
    return round(x), round(y)


def near(x, y, point):
    if abs(x - point[0]) < 10 and abs(y - point[1]) < 10:
        return True
    return False


def find_min_index(arr):
    min_index = 0
    for i in range(1, len(arr)):
        if arr[i] < arr[min_index]:
            min_index = i
    return min_index


def find_max_index(arr):
    max_index = 0
    for i in range(1, len(arr)):
        if arr[i] > arr[max_index]:
            max_index = i
    return max_index


def has_duplicate_x(points):
    x_set = set(points[:, 0])
    if len(x_set) < len(points):
        return True
    return False


class MplCanvas(FigureCanvasQTAgg):
    update_curve = pyqtSignal()
    counter = 0

    def __init__(self, target):
        self.fig = Figure(figsize=(400, 400))
        self.axes = self.fig.add_subplot(111)
        super(MplCanvas, self).__init__(self.fig)

        # 默认点集
        self.points = np.array([[0, 0], [255, 255]])  # 点集
        self.target = target

        # 当前正在拉取的点的索引
        self.curr_dragging_point_index = None

        self.draw_curve()

    def draw_curve(self):
        self.counter += 1

        x = np.array(self.points[:, 0])  # 点集的x坐标集合
        y = np.array(self.points[:, 1])  # 点集的y坐标集合
        xp = np.linspace(0, 255, 256)

        if len(self.points) < 4:
            # draw curve using polynomial curve fitting, but this is too is too "wavy"
            # if there are less than 4 points, we use this method
            z = np.polyfit(x, y, len(self.points) - 1)
            p = np.poly1d(z)
            yp = p(xp)
        else:
            # draw curve using cubic interpolation, which is more Ps-like
            f = interpolate.interp1d(x, y, 'cubic', fill_value="extrapolate")
            yp = f(xp)

        # crop yp into appropriate boundary
        for i in range(256):
            if yp[i] > 255:
                yp[i] = 255
            elif yp[i] < 0:
                yp[i] = 0

            # set boundary according to last point
            x_min = find_min_index(self.points[:, 0])
            x_max = find_max_index(self.points[:, 0])
            if i > self.points[x_max][0]:
                yp[i] = self.points[x_max][1]
            elif i < self.points[x_min][0]:
                yp[i] = self.points[x_min][1]

            # assign value to target
            self.target[i] = yp[i]

        # set up axes
        self.axes.clear()
        self.axes.set_xlim(0, 256)
        self.axes.set_ylim(0, 256)
        self.axes.autoscale(False, 'both')

        # plot curve
        self.axes.plot(x, y, '.', xp, yp)

        # do the following so that it refreshes
        if self.counter > 2:  # do not invoke this when initializing
            self.fig.canvas.draw()  # update plot
        self.fig.canvas.flush_events()

    def mousePressEvent(self, event):
        x, y = normalize(event.pos().x(), event.pos().y())
        # check if (x, y) is near any point in self.points
        for i in range(len(self.points)):
            if near(x, y, self.points[i]):
                self.curr_dragging_point_index = i
                break

        # create a new point
        if self.curr_dragging_point_index is None:
            self.points = np.append(self.points, [[x, y]], axis=0)
            self.curr_dragging_point_index = -1

    def mouseMoveEvent(self, event):
        if self.curr_dragging_point_index is not None:
            x, y = normalize(event.pos().x(), event.pos().y())
            if x > 255 or x < 0 or y > 255 or y < 0:
                return
            self.points[self.curr_dragging_point_index] = x, y

            # if two points form a vertical line
            if has_duplicate_x(self.points):
                self.points = np.delete(self.points, self.curr_dragging_point_index, axis=0)

                # stop drawing by mimicking a mouseReleaseEvent
                mouse_release_event = QtGui.QMouseEvent(
                    QtCore.QEvent.MouseButtonRelease,
                    self.cursor().pos(),
                    QtCore.Qt.LeftButton,
                    QtCore.Qt.LeftButton,
                    QtCore.Qt.NoModifier,
                )
                QtCore.QCoreApplication.postEvent(self, mouse_release_event)

            self.draw_curve()

    def mouseReleaseEvent(self, event):
        if self.curr_dragging_point_index is not None:
            x, y = normalize(event.pos().x(), event.pos().y())
            if x > 255 or x < 0 or y > 255 or y < 0:
                self.curr_dragging_point_index = None
                return
            self.points[self.curr_dragging_point_index] = x, y

            # if two points form a vertical line
            if has_duplicate_x(self.points):
                self.points = np.delete(self.points, self.curr_dragging_point_index, axis=0)
                # self.points[self.curr_dragging_point_index][0] -= 1

            self.curr_dragging_point_index = None
            self.draw_curve()
            self.update_curve.emit()


class MyCurveWidget(QWidget):
    update_window_image = pyqtSignal(str)  # 定义信号

    def __init__(self, mode, target):
        super().__init__()

        # set up widget
        self.ui = Ui_Curve()
        self.ui.setupUi(self)
        self.setFixedSize(QSize(400, 441))
        # set background color to white
        palette = QPalette()
        palette.setColor(QPalette.Background, Qt.white)
        self.setPalette(palette)
        # set up disability according to mode
        if mode == 'L':
            self.ui.label_1.setEnabled(False)
            self.ui.comboBox.setEnabled(False)

        self.target = target

        # matplotlib控件
        self.sc = MplCanvas(target)
        self.ui.verticalLayout.addWidget(self.sc)
        self.setLayout(self.ui.verticalLayout)

        # set up actions
        self.init_actions()
        self.update_curve()

        # 默认的通道是RGB, 当图像为灰度图时，RGB就意味着全局灰度
        self.channel = 'RGB'

    def init_actions(self):
        self.sc.update_curve.connect(self.update_curve)

    def update_curve(self):
        self.update_window_image.emit(self.ui.comboBox.currentText())


if __name__ == '__main__':
    app = QApplication(sys.argv)
    win = MyCurveWidget('RGB', target=np.zeros(256))
    win.show()
    sys.exit(app.exec_())
